package com.example.state;

import com.example.pojo.MyPattern;
import com.example.pojo.UserAction;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.Objects;

/**
 * @author lbq
 * @version 1.0
 * @date 2022/3/27 13:42
 */
public class BroadCastDemo {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(4);
        UserAction ac1 = new UserAction(1001L, "login");
        UserAction ac2 = new UserAction(1003L, "pay");
        UserAction ac3 = new UserAction(1002L, "car");
        UserAction ac4 = new UserAction(1001L, "logout");
        UserAction ac5 = new UserAction(1003L, "car");
        UserAction ac6 = new UserAction(1002L, "logout");

        final DataStreamSource<UserAction> data = env.fromElements(ac1, ac2, ac3, ac4, ac5, ac6);

        MyPattern myPattern1 = new MyPattern("login", "logout");
        MyPattern myPattern2 = new MyPattern("car", "logout");

        final DataStreamSource<MyPattern> patternData = env.fromElements(myPattern2);

        final KeyedStream<UserAction, Long> keyedStream = data.keyBy(UserAction::getUserId);

        final MapStateDescriptor<Void, MyPattern> bcDescriptor = new MapStateDescriptor<>("patterns", Types.VOID, Types.POJO(MyPattern.class));
        final BroadcastStream<MyPattern> broadcast = patternData.broadcast(bcDescriptor);

        keyedStream.connect(broadcast).process(new PatternEvaluator()).print();

        env.execute();
    }

    public static class PatternEvaluator extends KeyedBroadcastProcessFunction<Long, UserAction, MyPattern, Tuple2<Long, MyPattern>> {

        private ValueState<String> preActionState;

        @Override
        public void open(Configuration parameters) throws Exception {
            preActionState = getRuntimeContext().getState(new ValueStateDescriptor<String>("preAction", Types.STRING));
        }

        @Override
        public void processElement(UserAction value, ReadOnlyContext ctx, Collector<Tuple2<Long, MyPattern>> out) throws Exception {
            final ReadOnlyBroadcastState<Void, MyPattern> patterns = ctx.getBroadcastState(new MapStateDescriptor<>("patterns", Types.VOID, Types.POJO(MyPattern.class)));
            final MyPattern myPattern = patterns.get(null);
            final String preAction = preActionState.value();
            if (myPattern != null && preAction != null) {
                if (Objects.equals(preAction, myPattern.getAc1()) && Objects.equals(value.getAction(), myPattern.getAc2())) {
                    out.collect(new Tuple2<>(ctx.getCurrentKey(), myPattern));
                }
            }
            preActionState.update(value.getAction());
        }

        @Override
        public void processBroadcastElement(MyPattern value, Context ctx, Collector<Tuple2<Long, MyPattern>> out) throws Exception {
            final BroadcastState<Void, MyPattern> bcState = ctx.getBroadcastState(new MapStateDescriptor<>("patterns", Types.VOID, Types.POJO(MyPattern.class)));
            bcState.put(null, value);
        }
    }
}
